import os
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image

class MedicalImageDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.image_files = sorted(os.listdir(os.path.join(root_dir, 'images')))
        self.mask_files = sorted(os.listdir(os.path.join(root_dir, 'masks')))

    def __len__(self):
        return len(self.image_files)

    def __getitem__(self, idx):
        image_path = os.path.join(self.root_dir, 'images', self.image_files[idx])
        mask_path = os.path.join(self.root_dir, 'masks', self.mask_files[idx])
        image = Image.open(image_path).convert('RGB')
        mask = Image.open(mask_path).convert('L')

        if self.transform:
            image = self.transform(image)
            # 对mask应用相同的尺寸变换，使用最近邻插值保持整数标签
            mask = transforms.Resize((256, 256), interpolation=transforms.InterpolationMode.NEAREST)(mask)
            mask = transforms.ToTensor()(mask)

        return image, mask

# 定义数据预处理
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载本地数据集 1
# root_train = './data/eyeImages/train'
# root_test = './data/eyeImages/test'

# 加载本地数据集 2
root_train = './data/chestSegImages/train'
root_test = './data/chestSegImages/test'


train_dataset = MedicalImageDataset(root_train, transform=transform)
test_dataset = MedicalImageDataset(root_test, transform=transform)

# 创建数据加载器, shuffle=True 数据会在每个 epoch 开始时被随机打乱,训练时候使用
# 每个 epoch 包含多个 batch，具体数量取决于数据集大小和 batch_size。
# 例如，数据集有 1000 个样本，batch_size=100，则每个 epoch 包含 10 个 batch。
train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=4, shuffle=False)